"""
模型融合模块
结合XGBoost和LSTM模型的预测结果
"""

import numpy as np
import pandas as pd
from sklearn.ensemble import VotingClassifier
from sklearn.linear_model import LogisticRegression
import warnings
warnings.filterwarnings('ignore')

class ModelFusion:
    """模型融合器"""
    
    def __init__(self, xgboost_model, lstm_model):
        self.xgboost_model = xgboost_model
        self.lstm_model = lstm_model
        self.fusion_weights = {
            'xgboost': 0.6,  # XGBoost权重
            'lstm': 0.4      # LSTM权重
        }
        
    def predict(self, data):
        """
        融合预测
        
        Args:
            data (dict or DataFrame): 输入数据
            
        Returns:
            dict: 融合后的预测结果
        """
        try:
            # 获取XGBoost预测结果
            xgb_result = self.xgboost_model.predict(data)
            
            # 获取LSTM预测结果
            lstm_result = self.lstm_model.predict(data)
            
            # 融合预测结果
            fused_result = self._fuse_predictions(xgb_result, lstm_result)
            
            # 添加时间预测
            fused_result['time_predictions'] = self._predict_time_horizons(fused_result)
            
            return fused_result
            
        except Exception as e:
            print(f"模型融合预测错误: {e}")
            return self._generate_default_prediction()
    
    def _fuse_predictions(self, xgb_result, lstm_result):
        """融合XGBoost和LSTM的预测结果"""
        try:
            # 提取概率值
            xgb_ht_prob = xgb_result['hypertension_risk']['probability']
            xgb_dm_prob = xgb_result['diabetes_risk']['probability']
            
            # LSTM结果处理（如果有预测值）
            lstm_ht_prob = 0.3  # 默认值
            lstm_dm_prob = 0.2  # 默认值
            
            if 'predicted_values' in lstm_result:
                # 基于LSTM预测的生理指标计算风险
                pred_values = lstm_result['predicted_values']
                lstm_ht_prob = self._calculate_risk_from_physiological(pred_values, 'hypertension')
                lstm_dm_prob = self._calculate_risk_from_physiological(pred_values, 'diabetes')
            
            # 加权融合
            fused_ht_prob = (self.fusion_weights['xgboost'] * xgb_ht_prob + 
                           self.fusion_weights['lstm'] * lstm_ht_prob)
            
            fused_dm_prob = (self.fusion_weights['xgboost'] * xgb_dm_prob + 
                           self.fusion_weights['lstm'] * lstm_dm_prob)
            
            # 计算综合风险
            overall_risk = (fused_ht_prob + fused_dm_prob) / 2
            
            # 确定风险等级
            ht_risk_level = self._determine_risk_level(fused_ht_prob)
            dm_risk_level = self._determine_risk_level(fused_dm_prob)
            overall_risk_level = self._determine_risk_level(overall_risk)
            
            # 计算预测置信度
            confidence = self._calculate_confidence(xgb_result, lstm_result)
            
            return {
                'hypertension_risk': {
                    'probability': float(fused_ht_prob),
                    'risk_level': ht_risk_level,
                    'xgb_contribution': float(self.fusion_weights['xgboost'] * xgb_ht_prob),
                    'lstm_contribution': float(self.fusion_weights['lstm'] * lstm_ht_prob)
                },
                'diabetes_risk': {
                    'probability': float(fused_dm_prob),
                    'risk_level': dm_risk_level,
                    'xgb_contribution': float(self.fusion_weights['xgboost'] * xgb_dm_prob),
                    'lstm_contribution': float(self.fusion_weights['lstm'] * lstm_dm_prob)
                },
                'overall_risk': {
                    'probability': float(overall_risk),
                    'risk_level': overall_risk_level
                },
                'prediction_confidence': confidence,
                'model_contributions': {
                    'xgboost_weight': self.fusion_weights['xgboost'],
                    'lstm_weight': self.fusion_weights['lstm']
                },
                'lstm_insights': lstm_result.get('trend_analysis', {}),
                'timestamp': pd.Timestamp.now().isoformat()
            }
            
        except Exception as e:
            print(f"预测融合错误: {e}")
            return self._generate_default_prediction()
    
    def _calculate_risk_from_physiological(self, pred_values, disease_type):
        """基于生理指标预测值计算疾病风险"""
        try:
            systolic_bp = pred_values.get('systolic_bp', 120)
            diastolic_bp = pred_values.get('diastolic_bp', 80)
            blood_glucose = pred_values.get('blood_glucose', 100)
            
            if disease_type == 'hypertension':
                # 基于血压值计算高血压风险
                if systolic_bp >= 140 or diastolic_bp >= 90:
                    risk = 0.8
                elif systolic_bp >= 130 or diastolic_bp >= 85:
                    risk = 0.6
                elif systolic_bp >= 120 or diastolic_bp >= 80:
                    risk = 0.4
                else:
                    risk = 0.2
            elif disease_type == 'diabetes':
                # 基于血糖值计算糖尿病风险
                if blood_glucose >= 126:
                    risk = 0.8
                elif blood_glucose >= 110:
                    risk = 0.6
                elif blood_glucose >= 100:
                    risk = 0.4
                else:
                    risk = 0.2
            else:
                risk = 0.3
            
            return risk
            
        except Exception as e:
            print(f"生理指标风险计算错误: {e}")
            return 0.3
    
    def _determine_risk_level(self, probability):
        """确定风险等级"""
        if probability < 0.3:
            return 'low'
        elif probability < 0.6:
            return 'medium'
        else:
            return 'high'
    
    def _calculate_confidence(self, xgb_result, lstm_result):
        """计算预测置信度"""
        try:
            xgb_confidence = xgb_result.get('prediction_confidence', 0.8)
            lstm_confidence = lstm_result.get('prediction_confidence', 0.8)
            
            # 加权平均置信度
            confidence = (self.fusion_weights['xgboost'] * xgb_confidence + 
                         self.fusion_weights['lstm'] * lstm_confidence)
            
            return float(confidence)
            
        except Exception as e:
            print(f"置信度计算错误: {e}")
            return 0.8
    
    def _predict_time_horizons(self, prediction_result):
        """预测不同时间跨度的风险"""
        try:
            ht_prob = prediction_result['hypertension_risk']['probability']
            dm_prob = prediction_result['diabetes_risk']['probability']
            
            # 1年风险（相对稳定）
            ht_1y = ht_prob * 0.3
            dm_1y = dm_prob * 0.2
            
            # 3年风险（中等增长）
            ht_3y = ht_prob * 0.6
            dm_3y = dm_prob * 0.5
            
            # 5年风险（接近预测值）
            ht_5y = ht_prob * 0.9
            dm_5y = dm_prob * 0.8
            
            return {
                '1_year': {
                    'hypertension_risk': float(ht_1y),
                    'diabetes_risk': float(dm_1y),
                    'overall_risk': float((ht_1y + dm_1y) / 2)
                },
                '3_years': {
                    'hypertension_risk': float(ht_3y),
                    'diabetes_risk': float(dm_3y),
                    'overall_risk': float((ht_3y + dm_3y) / 2)
                },
                '5_years': {
                    'hypertension_risk': float(ht_5y),
                    'diabetes_risk': float(dm_5y),
                    'overall_risk': float((ht_5y + dm_5y) / 2)
                }
            }
            
        except Exception as e:
            print(f"时间跨度预测错误: {e}")
            return {
                '1_year': {'hypertension_risk': 0.1, 'diabetes_risk': 0.1, 'overall_risk': 0.1},
                '3_years': {'hypertension_risk': 0.2, 'diabetes_risk': 0.2, 'overall_risk': 0.2},
                '5_years': {'hypertension_risk': 0.3, 'diabetes_risk': 0.3, 'overall_risk': 0.3}
            }
    
    def _generate_default_prediction(self):
        """生成默认预测结果"""
        return {
            'hypertension_risk': {
                'probability': 0.25,
                'risk_level': 'low',
                'xgb_contribution': 0.15,
                'lstm_contribution': 0.10
            },
            'diabetes_risk': {
                'probability': 0.20,
                'risk_level': 'low',
                'xgb_contribution': 0.12,
                'lstm_contribution': 0.08
            },
            'overall_risk': {
                'probability': 0.225,
                'risk_level': 'low'
            },
            'prediction_confidence': 0.75,
            'model_contributions': {
                'xgboost_weight': 0.6,
                'lstm_weight': 0.4
            },
            'lstm_insights': {
                'bp_trend': 'stable',
                'glucose_trend': 'stable'
            },
            'time_predictions': {
                '1_year': {'hypertension_risk': 0.08, 'diabetes_risk': 0.06, 'overall_risk': 0.07},
                '3_years': {'hypertension_risk': 0.15, 'diabetes_risk': 0.12, 'overall_risk': 0.135},
                '5_years': {'hypertension_risk': 0.23, 'diabetes_risk': 0.18, 'overall_risk': 0.205}
            },
            'timestamp': pd.Timestamp.now().isoformat()
        }
    
    def update_fusion_weights(self, xgb_weight, lstm_weight):
        """更新融合权重"""
        if 0 <= xgb_weight <= 1 and 0 <= lstm_weight <= 1 and abs(xgb_weight + lstm_weight - 1) < 0.01:
            self.fusion_weights['xgboost'] = xgb_weight
            self.fusion_weights['lstm'] = lstm_weight
            print(f"融合权重已更新: XGBoost={xgb_weight}, LSTM={lstm_weight}")
        else:
            print("权重必须为0-1之间的值且总和为1")
    
    def get_model_performance_metrics(self):
        """获取模型性能指标"""
        return {
            'fusion_weights': self.fusion_weights,
            'xgboost_status': 'trained' if self.xgboost_model.is_trained else 'not_trained',
            'lstm_status': 'trained' if self.lstm_model.is_trained else 'not_trained',
            'expected_accuracy': 0.85,
            'expected_auc': 0.88
        }
